package com.yangzhe.algorithm.c036;

// 求完全二叉树的节点个数
// 测试链接 : https://leetcode.cn/problems/count-complete-tree-nodes/
public class Code09_CountCompleteTreeNodes_LeetCode {

    // 不提交这个类
    public static class TreeNode {
        public int val;
        public TreeNode left;
        public TreeNode right;

        public TreeNode(int val) {
            this.val = val;
        }
    }

    public static void main(String[] args) {
        // 初始化根节点
        TreeNode root = new TreeNode(1);

        // 构建左子树
        root.left = new TreeNode(2);
        root.left.left = new TreeNode(4);
        root.left.right = new TreeNode(5);

        // 构建右子树
        root.right = new TreeNode(3);
        root.right.left = new TreeNode(6);
        root.right.right = null;  // 显式设置空节点

        System.out.println(countNodes(root));
    }

    /**
     * 思路：不论用什么遍历方式都是O(n)的复杂度。
     * 可以通过满二叉树计算节点个数的公式计算个数。判断子树是否为满二叉树，只要计算深度即可，极大减少的复杂度O(log2 n)
     * 设二叉树的高度为h，则满二叉树节点个数 max_nodes = 2 ^ h - 1
     * <p>
     * 左神设计的递归流程不好理解，这里修改一下：
     * 设递归方法 count(node)，左树高度 leftH，右树高度 rightH
     * 1. 判断当前节点：是否满二叉树 = leftH == rightH
     * 2. 是满二叉树，计算 2 ^ leftH - 1
     * 3. 不是满二叉树， + 1是当前节点：count(node.left) + count(node.right) + 1
     *
     * @param head
     * @return
     */
    public static int countNodes(TreeNode head) {
        return count(head);
    }

    public static int count(TreeNode head) {
        if (head == null) {
            return 0;
        }

        int leftH = 0;
        int rightH = 0;

        // 计算左树高度
        TreeNode cur = head;
        while (cur != null) {
            leftH++;
            cur = cur.left;
        }

        // 计算右树高度
        cur = head;
        while (cur != null) {
            rightH++;
            cur = cur.right;
        }

        if (leftH == rightH) {
            // 左树高度=右树高度则是满二叉树，使用公式计算高度
            // 这里隐含了base case，因为高度为1的时候，一定是满二叉树，这里计算结果就是1
            return (int) Math.pow(2, leftH) - 1;
        } else {
            // 否则递归左树和右树
            return count(head.left) + count(head.right) + 1;
        }
    }

}
